import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

public class UpdateMatrix {

    public int[][] updateMatrix1(int[][] mat) {
        List<int[]> list = new ArrayList<>();
        int[][] ans = new int[mat.length][mat[0].length];
        for (int i = 0; i < mat.length; i++) {
            for (int j = 0; j < mat[0].length; j++) {
                if (mat[i][j] == 1) {
                    list.add(new int[]{i, j});
                    ans[i][j] = Integer.MAX_VALUE;
                }
            }
        }
        for (int i = 0; i < mat.length; i++) {
            for (int j = 0; j < mat[0].length; j++) {
                if (mat[i][j] == 0) {
                    for (int k = 0; k < list.size(); k++) {
                        int[] temp = list.get(k);
                        ans[temp[0]][temp[1]] = Math.min(ans[temp[0]][temp[1]], Math.abs(temp[0] - i) + Math.abs(temp[1] - j));
                    }
                }
            }
        }
        return ans;
    }

    public int[][] updateMatrix(int[][] mat) {
        int[][] ans = new int[mat.length][mat[0].length];
        LinkedList<int[]> queue = new LinkedList<>();
        int[][] temp = {{0, 1}, {1, 0}, {-1, 0}, {0, -1}};
        for (int i = 0; i < mat.length; i++) {
            for (int j = 0; j < mat[0].length; j++) {
                if (mat[i][j] == 1) {
                    ans[i][j] = Integer.MAX_VALUE;
                } else queue.add(new int[]{i, j});
            }
        }
        while (!queue.isEmpty()) {
            int[] poll = queue.poll();
            for (int[] step : temp) {
                int x = poll[0] + step[0];
                int y = poll[1] + step[1];
                if (x >= 0 && x < mat.length && y >= 0 && y < mat[0].length
                        && ans[x][y] > ans[poll[0]][poll[1]] + 1) {
                    ans[x][y] = ans[poll[0]][poll[1]] + 1;
                    queue.add(new int[]{x, y});
                }
            }
        }
        return ans;
    }
}
